{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pygame\n",
    "import numpy as np\n",
    "from Game import GameEnvironment\n",
    "from model import QNetwork, get_network_input\n",
    "import torch\n",
    "import sys\n",
    "\n",
    "gridsize = 15 # 13\n",
    "framerate = 10\n",
    "block_size = 20\n",
    "\n",
    "model = QNetwork(input_dim=10, hidden_dim=20, output_dim=5)\n",
    "## model.load_state_dict(torch.load('./dir_chk_lr0.0001/Snake_40000'))\n",
    "model.load_state_dict(torch.load('./dir_chk_lr0.00001/Snake_60000'))\n",
    "\n",
    "board = GameEnvironment(gridsize, 0., -100., 100.)\n",
    "windowwidth = gridsize*block_size*2\n",
    "windowheight = gridsize*block_size\n",
    "\n",
    "pygame.init()\n",
    "win = pygame.display.set_mode((windowwidth, windowheight))\n",
    "pygame.display.set_caption(\"snake\")\n",
    "font = pygame.font.SysFont('arial', 18)\n",
    "clock = pygame.time.Clock()\n",
    "\n",
    "def drawboard(snake, apple):\n",
    "    win.fill((0,0,0))\n",
    "    for pos in snake.prevpos:\n",
    "        pygame.draw.rect(win, (0,255,0), (pos[0]*block_size, pos[1]*block_size, block_size, block_size))\n",
    "    pygame.draw.rect(win, (255, 0, 0), (apple.pos[0]*block_size, apple.pos[1]*block_size, block_size, block_size))\n",
    "\n",
    "runGame = True\n",
    "\n",
    "prev_len_of_snake = 0\n",
    "\n",
    "while runGame:\n",
    "    clock.tick(framerate)\n",
    "    \n",
    "    state_0 = get_network_input(board.snake, board.apple)\n",
    "    state = model(state_0)\n",
    "    \n",
    "    action = torch.argmax(state)\n",
    "    \n",
    "    reward, done, len_of_snake = board.update_boardstate(action)\n",
    "    drawboard(board.snake, board.apple)\n",
    "    \n",
    "    lensnaketext     = font.render('          LEN OF SNAKE: ' + str(len_of_snake), False, (255, 255, 255))\n",
    "    rewardtext       = font.render('          REWARD: ' + str(int(reward)), False, (255, 255, 255))\n",
    "    prevlensnaketext = font.render('          LEN OF PREVIOUS SNAKE: ' + str(prev_len_of_snake), False, (255, 255, 255))\n",
    "    \n",
    "    win.blit(lensnaketext, (windowwidth//2, 40))\n",
    "    win.blit(rewardtext, (windowwidth//2, 80))\n",
    "    win.blit(prevlensnaketext, (windowwidth//2, 120))\n",
    "\n",
    "    \n",
    "    for event in pygame.event.get():\n",
    "        if event.type==pygame.KEYDOWN and event.key == pygame.K_ESCAPE:   \n",
    "            runGame = False\n",
    "    \n",
    "    keys = pygame.key.get_pressed()\n",
    "    if keys[pygame.K_r]:\n",
    "        paused = True\n",
    "        while paused == True:\n",
    "            clock.tick(10)\n",
    "            pygame.event.pump()\n",
    "            for event in pygame.event.get():\n",
    "                if event.type == pygame.KEYDOWN:\n",
    "                    paused = False\n",
    "    \n",
    "    pygame.display.update()\n",
    "    \n",
    "    if board.game_over == True:\n",
    "        prev_len_of_snake = len_of_snake\n",
    "        board.resetgame()\n",
    "\n",
    "pygame.quit()            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
